///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/ham/integral_transform.cpp
///
/// \brief   implementation file for functions involved in integral transformations
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<cmath>
#include<cassert>
#include<complex>
#include<vector>

#include<formic/ham/integral_transform.h>
#include<formic/ham/tei.h>
#include<formic/exception.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/lapack/interface.h>
#include<formic/timing/timing.h>
#include<formic/output/output.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   transforms the one electron integrals into a new orbital basis
///
/// \param[in]   n          the number of orbitals in the basis
/// \param[in]   new_orbs   the new orbital coefficients, size n*n
/// \param[in]   old_oei    the initial one electron integrals, size n*n
/// \param[out]  new_oei    the transformed one electron integrals, size n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::transform_oei(const int n,
                                             const S * const new_orbs,
                                             const S * const old_oei,
                                                   S * const new_oei)
{

  // create complex conjugate of orbital coefficients
  std::vector<S> ccnj_vec(n*n, formic::zero(S()));
  S * const ccnj = &ccnj_vec.at(0);
  for (int p = 0; p < n; p++)
  for (int i = 0; i < n; i++)
    ccnj[i+p*n] = formic::conj(new_orbs[i+p*n]);

  // get workspace
  std::vector<S> work_vec(n*n, formic::zero(S()));
  S * const work = &work_vec.at(0);

  // transform the one electron integrals
  formic::xgemm('N', 'N', n, n, n, formic::unity(S()), ccnj, n, old_oei,  n, formic::zero(S()), work,    n);
  formic::xgemm('N', 'T', n, n, n, formic::unity(S()), work, n, new_orbs, n, formic::zero(S()), new_oei, n);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   transforms the two electron integrals into a new orbital basis
///
/// \param[in]   n          the number of orbitals in the basis
/// \param[in]   new_orbs   the new orbital coefficients, size n*n
/// \param[in]   cnj_orbs   complex conjugate of the new orbital coefficients, size n*n
/// \param[in]   old_tei    the initial two electron integrals, size n*n*n*n
/// \param[out]  new_tei    the transformed two electron integrals, size n*n*n*n
/// \param[out]  work       work space, size n*n*n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::transform_dense_tei(const int n,
                                                   const S * const new_orbs,
                                                   const S * const cnj_orbs,
                                                   const S * const old_tei,
                                                         S * const new_tei,
                                                         S * const work)
{

//  // we will transform from initial orbitals (indexed by pi,qi,ri,si) to final orbitals (indexed by pf,qf,rf,sf)
//
//  // get n squared and n cubed
//  const int n2 = n*n;
//  const int n3 = n*n*n;
//
//  // get convenient pointers
//  S * const array1 = work;
//  S * const array2 = new_tei;
//
//  // convert (pi,qi,ri,si) to (pf,qi,ri,si)
//  formic::xgemm('N', 'N', n, n3, n, formic::unity(S()), cnj_orbs, n, old_tei, n, formic::zero(S()), array1, n);
//
//  // swap pf and qi
//  formic::tei_swap_first_two_indices(n, array1, array2);
//
//  // convert (qi,pf,ri,si) to (qf,pf,ri,si)
//  formic::xgemm('N', 'N', n, n3, n, formic::unity(S()), new_orbs, n, array2, n, formic::zero(S()), array1, n);
//
//  // swap pf and qf
//  formic::tei_swap_first_two_indices(n, array1, array2);
//
//  // convert (pf,qf,ri,si) to (pf,qf,ri,sf)
//  formic::xgemm('N', 'T', n3, n, n, formic::unity(S()), array2, n3, new_orbs, n, formic::zero(S()), array1, n3);
//
//  // swap ri and sf
//  formic::tei_swap_last_two_indices(n, array1, array2);
//
//  // convert (pf,qf,sf,ri) to (pf,qf,sf,rf)
//  formic::xgemm('N', 'T', n3, n, n, formic::unity(S()), array2, n3, cnj_orbs, n, formic::zero(S()), array1, n3);
//
//  // swap rf and sf
//  formic::tei_swap_last_two_indices(n, array1, array2);
//
//  // transformed integrals are now in new_tei in (pf,qf,rf,sf) order

//////////////////////////////////////
//  !!! old version of function !!!

  // transform two electron integrals index i
  formic::xgemm('N', 'N', n, n*n*n, n, formic::unity(S()), cnj_orbs, n, old_tei, n, formic::zero(S()), work, n);
  formic::tei_rotate_indices_left(n, work, new_tei);

  // transform two electron integrals index j
  formic::xgemm('N', 'N', n, n*n*n, n, formic::unity(S()), new_orbs, n, new_tei, n, formic::zero(S()), work, n);
  formic::tei_rotate_indices_left(n, work, new_tei);

  // transform two electron integrals index k
  formic::xgemm('N', 'N', n, n*n*n, n, formic::unity(S()), cnj_orbs, n, new_tei, n, formic::zero(S()), work, n);
  formic::tei_rotate_indices_left(n, work, new_tei);

  // transform two electron integrals index l
  formic::xgemm('N', 'N', n, n*n*n, n, formic::unity(S()), new_orbs, n, new_tei, n, formic::zero(S()), work, n);
  formic::tei_rotate_indices_left(n, work, new_tei);

//////////////////////////////////////

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   transforms the one and two electron integrals into a new orbital basis
///
/// \param[in]   n          the number of orbitals in the basis
/// \param[in]   old_oei    the initial one electron integrals, size n*n
/// \param[in]   old_tei    the initial two electron integrals, size n*n*n*n
/// \param[in]   new_orbs   the new orbital coefficients, size n*n
/// \param[out]  new_oei    the initial one electron integrals, size n*n
/// \param[out]  new_tei    the initial two electron integrals, size n*n*n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::transform_dense_oei_tei(const int n,
                                                       const S * const old_oei,
                                                       const S * const old_tei,
                                                       const S * const new_orbs,
                                                             S * const new_oei,
                                                             S * const new_tei)
{

//  // print for debugging
//  formic::of << boost::format("initial oei") << std::endl;
//  for (int i = 0; i < n*n; i++)
//    formic::of << boost::format("%12.4f") % old_oei[i] << std::endl;
//  formic::of << std::endl;
//  formic::of << boost::format("initial tei") << std::endl;
//  for (int i = 0; i < n*n*n*n; i++)
//    formic::of << boost::format("%12.4f") % old_tei[i] << std::endl;
//  formic::of << std::endl;
//  formic::of << boost::format("new orbs") << std::endl;
//  for (int i = 0; i < n*n; i++)
//    formic::of << boost::format("%12.4f") % new_orbs[i] << std::endl;
//  formic::of << std::endl;

  // get some workspace
  std::vector<S> work_vec(n*n*n*n, formic::zero(S()));
  S * const work = &work_vec.at(0);

  // create conjugate of orbital coefficients
  std::vector<S> ccnj_vec(n*n, formic::zero(S()));
  S * const ccnj = &ccnj_vec.at(0);
  for (int p = 0; p < n; p++)
  for (int i = 0; i < n; i++)
    ccnj[i+p*n] = formic::conj(new_orbs[i+p*n]);

  // create conjugate transpose of orbital coefficients
  std::vector<S> ctrn_vec(n*n, formic::zero(S()));
  S * const ctrn = &ctrn_vec.at(0);
  for (int i = 0; i < n; i++)
  for (int p = 0; p < n; p++)
    ctrn[p+i*n] = formic::conj(new_orbs[i+p*n]);

  // transform one electron integrals
  formic::transform_oei(n, new_orbs, old_oei, new_oei);
  //formic::xgemm('N', 'N', n, n, n, formic::unity(S()), new_orbs, n, old_oei, n, formic::zero(S()), work, n);
  //formic::xgemm('N', 'N', n, n, n, formic::unity(S()), work, n, ctrn, n, formic::zero(S()), new_oei, n);

  // transform two electron integrals
  formic::transform_dense_tei(n, new_orbs, ccnj, old_tei, new_tei, work);

//  // print for debugging
//  formic::of << boost::format("final oei") << std::endl;
//  for (int i = 0; i < n*n; i++)
//    formic::of << boost::format("%12.4f") % new_oei[i] << std::endl;
//  formic::of << std::endl;
//  formic::of << boost::format("final tei") << std::endl;
//  for (int i = 0; i < n*n*n*n; i++)
//    formic::of << boost::format("%12.4f") % new_tei[i] << std::endl;
//  formic::of << std::endl;

}

template void formic::transform_dense_oei_tei(const int n,
                                              const double * const old_oei,
                                              const double * const old_tei,
                                              const double * const new_orbs,
                                                    double * const new_oei,
                                                    double * const new_tei);
template void formic::transform_dense_oei_tei(const int n,
                                              const std::complex<double> * const old_oei,
                                              const std::complex<double> * const old_tei,
                                              const std::complex<double> * const new_orbs,
                                                    std::complex<double> * const new_oei,
                                                    std::complex<double> * const new_tei);
